import os
import sys
sys.path.append(os.path.abspath(os.path.dirname(__file__) + '/' + '../..'))

import matplotlib as mpl

mpl.use('pgf')
pgf_with_latex = {                     
    "pgf.texsystem": "xelatex",        
    "pgf.rcfonts": False,
    "text.usetex": True,                
    "font.family": "Times New Roman",
    "pgf.preamble": [
        r"\usepackage{fontspec}",    
        r"\setmainfont{Times New Roman}",        
        r"\usepackage{unicode-math}",
        r"\setmathfont{XITS Math}",
        ]
    }    
mpl.rcParams.update(pgf_with_latex)

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

from python.tools import (
    clean_folder,
    plot_IRF,
    plot_rolling_window_estimates,
    construct_state_dep_table
)

################
## Parameters ##
################

input_folder = './get_graphs/input'
output_folder = './get_graphs/output/graphs'

###############
## Plot IRFs ##
###############

clean_folder(output_folder)

# Impulse response of inflation
df = pd.read_csv('{}/IRF_inflation.txt'.format(input_folder), sep = '\t')
df = df.loc[df['lag'] <= 12, ]
plot_IRF(df, '{}/IRF_inflation'.format(output_folder), location_legend = "upper right", 
  y_label = 'Impulse Response Function', x_label = 'Time (Quarters)')

# Composite coefficients from consensus data
df = pd.read_csv('{}/IRF_SPF_consensus.txt'.format(input_folder), sep = '\t')
df = df.loc[df['lag'] <= 12, ]
plot_IRF(df, '{}/bias_coef_consensus'.format(output_folder), 
  bias_coefs = True, ylims = (-0.6, 0.35),
  y_label = 'Bias Coefficients', x_label = 'Lag (Quarters)')

# Composite coefficients from consensus data using maximum likelihood
df = pd.read_csv('{}/IRF_SPF_consensus_MA.txt'.format(input_folder), sep = '\t')
df = df.loc[df['lag'] <= 12, ]
plot_IRF(df, '{}/bias_coef_consensus_MA'.format(output_folder), bias_coefs = True, ylims = (-0.6, 0.35),
  y_label = 'Bias Coefficients', x_label = 'Lag (Quarters)')

# Composite coefficients from individual-level data
df = pd.read_csv('{}/IRF_SPF_individual.txt'.format(input_folder), sep = '\t')
df = df.loc[df['lag'] <= 12, ]
plot_IRF(df, '{}/bias_coef_individual'.format(output_folder), 
  bias_coefs = True, ylims = (-0.6, 0.35),
  y_label = 'Bias Coefficients', x_label = 'Lag (Quarters)')

# Composite coefficients from individual-level data (winsorized)
df = pd.read_csv('{}/IRF_SPF_individual_winsorized.txt'.format(input_folder), sep = '\t')
df = df.loc[df['lag'] <= 12, ]
plot_IRF(df, '{}/bias_coef_individual_winsorized'.format(output_folder), 
  bias_coefs = True, ylims = (-0.6, 0.35),
  y_label = 'Bias Coefficients', x_label = 'Lag (Quarters)')

# Idiosyncratic bias coefficients
df = pd.read_csv('{}/IRF_idiosyncratic.txt'.format(input_folder), sep = '\t')
mask = (df['lag'] <= 13)
df = df.loc[mask, ]
df['lag'] = (df['lag'] - 1)
plot_IRF(df, '{}/bias_coef_idiosyncratic'.format(output_folder), bias_coefs = True,
  location_legend = 'upper right', ylims = (-0.30, 1.15),
  y_label = 'Bias Coefficients', x_label = 'Lag (Quarters)')

# Idiosyncratic bias coefficients
df = pd.read_csv('{}/IRF_idiosyncratic_winsorized.txt'.format(input_folder), sep = '\t')
mask = (df['lag'] <= 13)
df = df.loc[mask, ]
df['lag'] = (df['lag'] - 1)
plot_IRF(df, '{}/bias_coef_idiosyncratic_winsorized'.format(output_folder), bias_coefs = True,
  location_legend = 'upper right', ylims = (-0.30, 1.15),
  y_label = 'Bias Coefficients', x_label = 'Lag (Quarters)')

# Idiosyncratic bias coefficients: JK shocks
df = pd.read_csv('{}/IRF_idiosyncratic_JK_shocks.txt'.format(input_folder), sep = '\t')
mask = (df['lag'] <= 13)
df = df.loc[mask, ]
df['lag'] = (df['lag'] - 1)
plot_IRF(df, '{}/bias_coef_idiosyncratic_JK_shocks'.format(output_folder), bias_coefs = True,
  location_legend = 'upper right', ylims = (-0.30, 1.15),
  y_label = 'Bias Coefficients', x_label = 'Lag (Quarters)')

# Aggregate bias coefficients using Romer-Romer shocks
df = pd.read_csv('{}/IRF_SPF_Romer_Romer.txt'.format(input_folder), sep = '\t')
df = df.loc[df['lag'] <= 20, ]
df['IRF'] = (-1) * df['IRF'] # Since the shock is disinflationary
plot_IRF(df, '{}/bias_coef_Romer_Romer'.format(output_folder), location_legend = "upper right", 
  bias_coefs = True, y_label = 'Bias Coefficients', x_label = 'Lag (Quarters)')

# Response of inflation to Romer-Romer shocks
df = pd.read_csv('{}/IRF_inflation_Romer_Romer.txt'.format(input_folder), sep = '\t')
df = df.loc[df['lag'] <= 20, ]
plot_IRF(df, '{}/IRF_inflation_Romer_Romer'.format(output_folder), location_legend = "upper right",
  y_label = 'Impulse Response Function', x_label = 'Time (Quarters)')

# Graph of bias coefficients across different horizons
max_L = 4 # Number of bias coefficients

# Collect data: Loop across forecast horizons (hh)
for hh in range(1, 4 + 1):
    df_temp = pd.read_csv('{}/IRF_SPF_consensus_{}Q_ahead.txt'.format(input_folder, hh), sep = '\t')
    df_temp['bias_coef'] = -df_temp['IRF']
    mask = (df_temp['lag'] >= hh) & (df_temp['lag'] <= max_L + hh - 1)
    df_temp = df_temp.loc[mask, ].copy()
    df_temp = df_temp[['lag', 'bias_coef', 'se']]
    df_temp['horizon'] = hh
    df_temp['lag'] = (df_temp['lag'] + 1) - hh
    if hh == 1:
        df = df_temp
    else:
        df = df.append(df_temp)

# Get graph
fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, sharex=True, sharey=True)

mask = df['lag'] == 1
ax1.bar(df.loc[mask, 'horizon'], 
        df.loc[mask, 'bias_coef'],
        yerr = df.loc[mask, 'se'],
        alpha = 0.8,
        error_kw=dict(ecolor='black', lw=1.5, alpha = 0.8))
ax1.set_title('Lag $=1$')
ax1.set_ylim(-0.6, 0.1)
ax1.set_ylabel('Bias Coefficients', fontsize = 12)

mask = df['lag'] == 2
ax2.bar(df.loc[mask, 'horizon'], 
        df.loc[mask, 'bias_coef'],
        yerr = df.loc[mask, 'se'],
        alpha = 0.8,
        error_kw=dict(ecolor='black', lw=1.5, alpha = 0.8))
ax2.set_title('Lag $= 2$')
ax2.set_ylim(-0.6, 0.1)

mask = df['lag'] == 3
ax3.bar(df.loc[mask, 'horizon'], 
        df.loc[mask, 'bias_coef'],
        yerr = df.loc[mask, 'se'],
        alpha = 0.8,
        error_kw=dict(ecolor='black', lw=1.5, alpha = 0.8))
ax3.set_title('Lag $= 3$')
ax3.set_ylim(-0.6, 0.1)
ax3.set_xlabel('Horizon (Quarters)', fontsize = 12)
ax3.set_ylabel('Bias Coefficients', fontsize = 12)

mask = df['lag'] == 4
ax4.bar(df.loc[mask, 'horizon'], 
        df.loc[mask, 'bias_coef'],
        yerr = df.loc[mask, 'se'],
        alpha = 0.8,
        error_kw=dict(ecolor='black', lw=1.5, alpha = 0.8))
ax4.set_title('Lag $= 4$')
ax4.set_ylim(-0.6, 0.1)
ax4.set_xlabel('Horizon (Quarters)', fontsize = 12)

fig.tight_layout()
plt.savefig('{}/horizons.png'.format(output_folder), dpi = 600, bbox_inches='tight')
plt.savefig('{}/horizons.pgf'.format(output_folder), dpi = 600, bbox_inches='tight')
df.to_csv('{}/horizons.csv'.format(output_folder), index = False)

# Get rolling window estimates for composite bias coefficients
df = pd.read_csv('{}/IRF_SPF_consensus_rolling.csv'.format(input_folder))
mask = ~np.isnan(df['IRF'])
df = df.loc[mask, ]
df['date'] = df['date'].map(lambda x: x.replace(':0', 'Q'))
df['date'] = pd.to_datetime(df['date'])
df.index = df['date'] 
del df['date']
df['IRF'] = -df['IRF']
plot_rolling_window_estimates(df, '{}/bias_coef_consensus_rolling'.format(output_folder), 
        y_label = 'Bias Coefficient', location_legend = 'upper left')

# Get rolling window estimates for actual inflation
df = pd.read_csv('{}/IRF_inflation_rolling.csv'.format(input_folder))
mask = ~np.isnan(df['IRF'])
df = df.loc[mask, ]
df['date'] = df['date'].map(lambda x: x.replace(':0', 'Q'))
df['date'] = pd.to_datetime(df['date'])
df.index = df['date'] 
del df['date']
plot_rolling_window_estimates(df, '{}/IRF_inflation_rolling'.format(output_folder), 
        y_label = 'Impulse Response', location_legend = 'lower left')

############
## Tables ##
############

# CG regressions

get_se = lambda df, var_name: np.sqrt(df.loc[df['Unnamed: 0'] == var_name, var_name].values[0])

# Collect results
df = pd.read_csv('{}/CG_coefs_consensus.txt'.format(input_folder), sep = '\t')
cons_coef = df['Rfor_Step2'].values[0]

df = pd.read_csv('{}/CG_cov_consensus.txt'.format(input_folder), sep = '\t')
cons_se = get_se(df, 'Rfor_Step2')

df = pd.read_csv('{}/CG_coefs_micro_no_FE.txt'.format(input_folder), sep = '\t')
indiv_coef_no_FE = df['Rfor_Step2'].values[0]

df = pd.read_csv('{}/CG_cov_micro_no_FE.txt'.format(input_folder), sep = '\t')
indiv_se_no_FE = get_se(df, 'Rfor_Step2')

df = pd.read_csv('{}/CG_coefs_micro_FE.txt'.format(input_folder), sep = '\t')
indiv_coef_FE = df['Rfor_Step2'].values[0]

df = pd.read_csv('{}/CG_cov_micro_FE.txt'.format(input_folder), sep = '\t')
indiv_se_FE = get_se(df, 'Rfor_Step2')

df = pd.read_csv('{}/CG_coefs_idio_no_FE.txt'.format(input_folder), sep = '\t')
idio_coef_no_FE = df['Rfor_Step2'].values[0]

df = pd.read_csv('{}/CG_cov_idio_no_FE.txt'.format(input_folder), sep = '\t')
idio_se_no_FE = get_se(df, 'Rfor_Step2')

df = pd.read_csv('{}/CG_coefs_idio_FE.txt'.format(input_folder), sep = '\t')
idio_coef_FE = df['Rfor_Step2'].values[0]

df = pd.read_csv('{}/CG_cov_idio_FE.txt'.format(input_folder), sep = '\t')
idio_se_FE = get_se(df, 'Rfor_Step2')

# Calculate implied weights on idiosyncratic shocks
weight_no_FE = (cons_coef - indiv_coef_no_FE) / (cons_coef - idio_coef_no_FE)
weight_FE = (cons_coef - indiv_coef_FE) / (cons_coef - idio_coef_FE)

# Save to LaTeX
latex_template = """
\\begin{{tabular}}{{lcc}}
\\toprule
 & (1) & (2) \\\\
\\midrule
$\\hat{{\\beta}}^\\text{{macro}}_\\text{{CG}}$ & {cons_coef:.2f} & -- \\\\
 & {{\\footnotesize ({cons_se:.2f})}} &  \\\\
$\\hat{{\\beta}}^\\text{{micro}}_\\text{{CG}}$ & {indiv_coef_no_FE} & {indiv_coef_FE} \\\\
 & {{\\footnotesize ({indiv_se_no_FE})}} & {{\\footnotesize ({indiv_se_FE})}} \\\\
$\\hat{{\\beta}}^\\text{{idio.}}_\\text{{CG}}$ & {idio_coef_no_FE} & {idio_coef_FE} \\\\
 & {{\\footnotesize ({idio_se_no_FE})}} & {{\\footnotesize ({idio_se_FE})}} \\\\
\\midrule
Weight on Idiosyncratic & {weight_no_FE} & {weight_FE} \\\\
\\midrule
Forecaster FE & & $\checkmark$ \\\\
\\bottomrule
\\end{{tabular}}
"""

latex_template = latex_template.format(indiv_coef_no_FE = round(indiv_coef_no_FE, 2),
                                       indiv_coef_FE = round(indiv_coef_FE, 2),
                                       indiv_se_no_FE = round(indiv_se_no_FE, 2),
                                       indiv_se_FE = round(indiv_se_FE, 2),
                                       cons_coef = round(cons_coef, 2),
                                       cons_se = round(cons_se, 2),
                                       idio_coef_no_FE = round(idio_coef_no_FE, 2),
                                       idio_coef_FE = round(idio_coef_FE, 2),
                                       idio_se_no_FE = round(idio_se_no_FE, 2),
                                       idio_se_FE = round(idio_se_FE, 2),
                                       weight_no_FE = round(weight_no_FE, 2),
                                       weight_FE = round(weight_FE, 2))

with open('{}/CG_regressions.tex'.format(output_folder), "w") as text_file:
    text_file.write(latex_template)

# Reduced-form regressions: Example

# Collect results
df = pd.read_csv('{}/ex_CG_coefs_consensus.txt'.format(input_folder), sep = '\t')
CG_coef_cons = df['Rfor_Step2'].values[0]

df = pd.read_csv('{}/ex_CG_cov_consensus.txt'.format(input_folder), sep = '\t')
CG_se_cons = get_se(df, 'Rfor_Step2')

df = pd.read_csv('{}/ex_CG_coefs_micro.txt'.format(input_folder), sep = '\t')
CG_coef_micro = df['Rfor_Step2'].values[0]

df = pd.read_csv('{}/ex_CG_cov_micro.txt'.format(input_folder), sep = '\t')
CG_se_micro = get_se(df, 'Rfor_Step2')

df = pd.read_csv('{}/ex_autocorr_coefs_consensus.txt'.format(input_folder), sep = '\t')
autocorr_coef_cons = df['L1_Efor_Step2'].values[0]

df = pd.read_csv('{}/ex_autocorr_cov_consensus.txt'.format(input_folder), sep = '\t')
autocorr_se_cons = get_se(df, 'L1_Efor_Step2')

df = pd.read_csv('{}/ex_autocorr_coefs_micro.txt'.format(input_folder), sep = '\t')
autocorr_coef_micro = df['L1_Efor_Step2'].values[0]

df = pd.read_csv('{}/ex_autocorr_cov_micro.txt'.format(input_folder), sep = '\t')
autocorr_se_micro = get_se(df, 'L1_Efor_Step2')

df = pd.read_csv('{}/ex_extrap_coefs_consensus.txt'.format(input_folder), sep = '\t')
extrap_coef_cons = df['L1_Realiz1'].values[0]

df = pd.read_csv('{}/ex_extrap_cov_consensus.txt'.format(input_folder), sep = '\t')
extrap_se_cons = get_se(df, 'L1_Realiz1')

df = pd.read_csv('{}/ex_extrap_coefs_micro.txt'.format(input_folder), sep = '\t')
extrap_coef_micro = df['L1_Realiz1'].values[0]

df = pd.read_csv('{}/ex_extrap_cov_micro.txt'.format(input_folder), sep = '\t')
extrap_se_micro = get_se(df, 'L1_Realiz1')

df = pd.read_csv('{}/N_micro.txt'.format(input_folder), sep = '\t')
N_micro = df['c1'].values[0]

df = pd.read_csv('{}/N_cons.txt'.format(input_folder), sep = '\t')
N_cons = df['c1'].values[0]

# Save to LaTeX
latex_template = """
\\begin{{tabular}}{{lcc}}
\\toprule
 & Individual & Consensus \\\\
\\midrule
Coibion-Gorodnichenko & {CG_coef_micro:.2f} & {CG_coef_cons:.2f} \\\\
 & {{\\footnotesize ({CG_se_micro:.2f})}} & {{\\footnotesize ({CG_se_cons:.2f})}} \\\\
Extrapolation & {extrap_coef_micro:.2f} & {extrap_coef_cons:.2f} \\\\
 & {{\\footnotesize ({extrap_se_micro:.2f})}} & {{\\footnotesize ({extrap_se_cons:.2f})}} \\\\
Autocorrelation & {autocorr_coef_micro:.2f} & {autocorr_coef_cons:.2f} \\\\
 & {{\\footnotesize ({autocorr_se_micro:.2f})}} & {{\\footnotesize ({autocorr_se_cons:.2f})}} \\\\
\\midrule
$N$ & {N_micro} & {N_cons} \\\\
\\bottomrule
\\end{{tabular}}
"""

latex_template = latex_template.format(extrap_coef_micro = extrap_coef_micro,
                                       extrap_coef_cons = extrap_coef_cons,
                                       extrap_se_micro = extrap_se_micro,
                                       extrap_se_cons = extrap_se_cons,
                                       autocorr_coef_micro = autocorr_coef_micro,
                                       autocorr_coef_cons = autocorr_coef_cons,
                                       autocorr_se_micro = autocorr_se_micro,
                                       autocorr_se_cons = autocorr_se_cons,
                                       CG_coef_micro = CG_coef_micro,
                                       CG_coef_cons = CG_coef_cons,
                                       CG_se_micro = CG_se_micro,
                                       CG_se_cons = CG_se_cons,
                                       N_micro = N_micro,
                                       N_cons = N_cons)

with open('{}/reduced_form.tex'.format(output_folder), "w") as text_file:
    text_file.write(latex_template)

###########################################
## Get tables on state-dependent effects ##
###########################################

# For bias coefficients
IRF_baseline = pd.read_csv('{}/IRF_SPF_baseline.txt'.format(input_folder), sep = '\t')
IRF_high_infl = pd.read_csv('{}/IRF_SPF_high_inflation.txt'.format(input_folder), sep = '\t')
IRF_recession = pd.read_csv('{}/IRF_SPF_NBER_recession.txt'.format(input_folder), sep = '\t')
IRF_great_moderation = pd.read_csv('{}/IRF_SPF_great_moderation.txt'.format(input_folder), sep = '\t')
IRF_large_fe = pd.read_csv('{}/IRF_SPF_large_fe.txt'.format(input_folder), sep = '\t')

IRFs = [IRF_baseline, IRF_high_infl, IRF_recession, IRF_great_moderation, IRF_large_fe]
IRF_names = ['Baseline', 'High Inflation', 'Recession', 'Great Moderation', 'Large FE']    
IRF_baseline['t_stat'] = np.nan

# Convert to bias coefficients
for IRF in IRFs:
    IRF['IRF'] *= (-1)
    IRF['t_stat'] *= (-1)
    
table = construct_state_dep_table(IRFs, IRF_names, max_K = 4)
with open('{}/state_dep_bias_coefs.tex'.format(output_folder), "w") as text_file:
    text_file.write(table)

# For actual inflation
IRF_baseline = pd.read_csv('{}/IRF_inflation_baseline.txt'.format(input_folder), sep = '\t')
IRF_high_infl = pd.read_csv('{}/IRF_inflation_high_inflation.txt'.format(input_folder), sep = '\t')
IRF_recession = pd.read_csv('{}/IRF_inflation_NBER_recession.txt'.format(input_folder), sep = '\t')
IRF_great_moderation = pd.read_csv('{}/IRF_inflation_great_moderation.txt'.format(input_folder), sep = '\t')
IRF_large_fe = pd.read_csv('{}/IRF_inflation_large_fe.txt'.format(input_folder), sep = '\t')

IRFs = [IRF_baseline, IRF_high_infl, IRF_recession, IRF_great_moderation, IRF_large_fe]
IRF_names = ['Baseline', 'High Inflation', 'Recession', 'Great Moderation', 'Large FE']    
IRF_baseline['t_stat'] = np.nan
    
table = construct_state_dep_table(IRFs, IRF_names, max_K = 4)
with open('{}/state_dep_inflation.tex'.format(output_folder), "w") as text_file:
    text_file.write(table)